# encodign=utf-8

import cv2
import onnxruntime
import numpy


# load
ori_img = cv2.imread("sample/ori.jpg")
ori_size = ori_img.shape[:2]  # h, w

# transform
trans_image = cv2.cvtColor(ori_img, cv2.COLOR_BGR2RGB)
trans_image = cv2.resize(trans_image, (512, 512), interpolation=cv2.INTER_CUBIC)
trans_image = trans_image / 255.0

# return
x = numpy.expand_dims(trans_image.astype(numpy.float32).transpose((2, 0, 1)), axis=0)

# infer
session = onnxruntime.InferenceSession("export_oct_deeplabv3.onnx")
inputs = {"input":  x}
out = session.run(None, inputs)
print(out[0].shape)

# decode
decode_array = numpy.argmax(out[0], axis=1)[0]
H, W = ori_size
img_arr = cv2.resize(decode_array.astype(numpy.uint8), (W, H), interpolation=cv2.INTER_NEAREST)
cv2.imwrite("raw.png", img_arr)

# draw color
colors = {
    1: (0, 0, 255),
    2: (0, 255, 0),
    3: (255, 0, 0),
    4: (0, 255, 255),
    5: (255, 0, 255)
}
masks = numpy.zeros((H, W, 3), dtype=numpy.uint8)
for index, color in colors.items():
    masks[img_arr == index] = color

cv2.imwrite("color.png", masks)